import faiss
import numpy as np

# 假设我们有一些 128 维的向量
nb, d = 10000, 128  # 10000 个向量，每个向量 128 维
xb = np.random.random((nb, d)).astype('float32')
# 创建一个 HNSW 索引（适用于 GPU 和 CPU）
index = faiss.IndexHNSWFlat(d, 16)  # d 是向量维度，16 是 M 参数（影响性能和准确性）
# 或者使用其他索引类型，如 IndexFlatL2（适用于 CPU）
# index = faiss.IndexFlatL2(d)

# 将向量添加到索引中
index.add(xb)
# 准备一些查询向量
nq = 5  # 查询向量的数量
xq = np.random.random((nq, d)).astype('float32')

# 搜索与每个查询向量最接近的 k 个向量
k = 4  # 返回最接近的 k 个向量
D, I = index.search(xq, k)  # D 是距离数组，I 是索引数组

# 打印结果
print("查询向量与最近邻的距离:")
print(D)
print("最近邻的索引:")
print(I)

